from common.model.vecmodel_paper_segment import PaperSegmentModel
from common.rpc.paperfile import paperfile_pb2_grpc
from common.rpc.stub import new_rpc_stub
from langchain_openai import ChatOpenAI
from minio import Minio
from pymilvus import MilvusClient
from langchain_community.embeddings import DashScopeEmbeddings

from app.jobs.pdfparser.internal.config import config

class ServiceContext:
    def __init__(self, c: config.Config):
        self.config = c

        milvus_client = MilvusClient(uri=c.milvus_conf.uri)
        self.paper_segment_model = PaperSegmentModel(milvus_client)
        self.paperfile_stub = new_rpc_stub(self.config.paperfile_stub_conf, paperfile_pb2_grpc.PaperFileStub)
        self.minio_client = Minio(endpoint=c.minio_conf.endpoint, 
                                  access_key=c.minio_conf.access_key,
                                  secret_key=c.minio_conf.secret_key,
                                  secure=False)
        self.embeddings = DashScopeEmbeddings(
            model=c.ali_embeddings_conf.model,
            dashscope_api_key=c.ali_embeddings_conf.api_key,
        )
        self.llm = ChatOpenAI(
            api_key=c.ali_llm_conf.api_key,
            base_url=c.ali_llm_conf.base_url,
            model=c.ali_llm_conf.model,
            # other params...
            extra_body={"enable_thinking": False}
        )

